{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\" \n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 315,
   "metadata": {},
   "outputs": [],
   "source": [
    "file1 = open('../../datasets/RCV1/llm_cluster_result/update_labelspace.txt', 'r')\n",
    "documents = file1.readlines()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 316,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Airline Industry\\n', 'Politics\\n', 'Sports\\n', 'Intelligence Agencies\\n']"
      ]
     },
     "execution_count": 316,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "documents[:4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 317,
   "metadata": {},
   "outputs": [],
   "source": [
    "file2 = open('../../datasets/RCV1/train_label.txt', 'r')\n",
    "labels = file2.readlines() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 318,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_label_space = []\n",
    "for doc in documents:\n",
    "    predict_label_space.append(doc.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 208,
   "metadata": {},
   "outputs": [],
   "source": [
    "# predict_label_space = []\n",
    "# for doc in documents:\n",
    "#     label_list = doc.strip().split(': ')[1]\n",
    "#     if not label_list in predict_label_space:\n",
    "#         predict_label_space.append(label_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 319,
   "metadata": {},
   "outputs": [],
   "source": [
    "file2 = open('../../datasets/RCV1-V2/rcv1.topics.hier.orig.txt', 'r')\n",
    "labels = file2.readlines() \n",
    "# file2 = open('../../datasets/AAPD/true_labels.txt', 'r')\n",
    "# labels = file2.readlines()   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 320,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_label = []\n",
    "for row in labels[1:]:\n",
    "    new_label = row.strip().split('child-description: ')[1].lower()\n",
    "    true_label.append(new_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 321,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['strategy/plans',\n",
       " 'legal/judicial',\n",
       " 'regulation/policy',\n",
       " 'share listings',\n",
       " 'performance']"
      ]
     },
     "execution_count": 321,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "true_label[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 282,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Information Retrieval',\n",
       " 'Methodology',\n",
       " 'Quantum Physics',\n",
       " 'Information Theory',\n",
       " 'Applications']"
      ]
     },
     "execution_count": 282,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label_dic = {}\n",
    "for row in labels:\n",
    "    key = row.strip().split(';')[0].split('.')[0]\n",
    "    #key + '.' + \n",
    "    label_dic[row.strip().split(';')[0]] = row.strip().split(';')[1][1:]\n",
    "true_label = []\n",
    "for key in label_dic:\n",
    "    if not label_dic[key] in true_label:\n",
    "        true_label.append(label_dic[key])\n",
    "true_label[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['money-fx', 'earn', 'acq', 'sugar', 'trade']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "true_label = []\n",
    "count = 0\n",
    "for row in labels[2000:3000]:\n",
    "    label_list = row.strip().split(' ')\n",
    "    if len(label_list) == 1:\n",
    "        count +=1\n",
    "    for label in label_list:\n",
    "        if not label in true_label:\n",
    "            true_label.append(label)\n",
    "true_label[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "852"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 295,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['cocoa', 'sorghum', 'oat', 'barley', 'corn']"
      ]
     },
     "execution_count": 295,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "true_label = []\n",
    "for row in labels:\n",
    "    label_list = row.strip().split(' ')\n",
    "    for label in label_list:\n",
    "        if not label in true_label:\n",
    "            true_label.append(label)\n",
    "true_label[:5]\n",
    "\n",
    "# true_label = []\n",
    "# for row in labels:\n",
    "#     true_label.append(row.strip().split('\\t')[1])\n",
    "# true_label[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 229,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Similarity: tensor([[0.4893]])\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer, util\n",
    "model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')\n",
    "\n",
    "query_embedding = model.encode('teethers')\n",
    "passage_embedding = model.encode(['baby_dental_care'])\n",
    "\n",
    "print(\"Similarity:\", util.dot_score(query_embedding, passage_embedding))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 322,
   "metadata": {},
   "outputs": [],
   "source": [
    "passage_embedding = model.encode(predict_label_space)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 323,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_matrix = np.empty((0,len(predict_label_space)))\n",
    "for i in range(len(true_label)):\n",
    "    query_embedding = model.encode(true_label[i])\n",
    "    sim_matrix = np.append(sim_matrix, util.dot_score(query_embedding, passage_embedding).numpy(), 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 324,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(103, 110)"
      ]
     },
     "execution_count": 324,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim_matrix.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 325,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['legal/judicial',\n",
       "  (0.2560056447982788, 'Technology/Internet'),\n",
       "  (0.25780072808265686, 'Operations'),\n",
       "  (0.25806641578674316, 'Protests/Strikes'),\n",
       "  (0.2604048550128937, 'Historical Events'),\n",
       "  (0.2673090100288391, 'Resignation'),\n",
       "  (0.2749345302581787, 'Terrorism'),\n",
       "  (0.27502572536468506, 'Sales/Revenue'),\n",
       "  (0.27564650774002075, 'Labor/Workforce'),\n",
       "  (0.27792787551879883, 'Auctions'),\n",
       "  (0.27962490916252136, 'Privatization'),\n",
       "  (0.2804669737815857, 'Geopolitics'),\n",
       "  (0.28190767765045166, 'Death'),\n",
       "  (0.2854553461074829, 'Environmental'),\n",
       "  (0.28918778896331787, 'European Union'),\n",
       "  (0.29928943514823914, 'Financial/Economic'),\n",
       "  (0.3115212321281433, 'Sports'),\n",
       "  (0.3132980465888977, 'Trade Regulations'),\n",
       "  (0.3178219795227051, 'Competition'),\n",
       "  (0.31876689195632935, 'Location'),\n",
       "  (0.3215720057487488, 'Education/Training'),\n",
       "  (0.3246457278728485, 'Intelligence Agencies'),\n",
       "  (0.32509419322013855, 'Military'),\n",
       "  (0.3338983654975891, 'Real Estate'),\n",
       "  (0.34357425570487976, 'Literature'),\n",
       "  (0.3485540747642517, 'Elections'),\n",
       "  (0.35694149136543274, 'Political Violence'),\n",
       "  (0.36028754711151123, 'Taxation'),\n",
       "  (0.36262935400009155, 'Conflict'),\n",
       "  (0.362663209438324, 'Religion'),\n",
       "  (0.37119609117507935, 'News/Media'),\n",
       "  (0.37882280349731445, 'Corruption'),\n",
       "  (0.3812379837036133, 'Government/Public Sector'),\n",
       "  (0.38202032446861267, 'Regulation'),\n",
       "  (0.38607972860336304, 'Bankruptcy'),\n",
       "  (0.396829217672348, 'Science/Research'),\n",
       "  (0.4446495473384857, 'Politics'),\n",
       "  (0.491844117641449, 'Crime'),\n",
       "  (0.5573071241378784, 'Legislation'),\n",
       "  (0.6152762174606323, 'Legal'),\n",
       "  (0.7028183937072754, 'Legal/Court Cases')],\n",
       " ['regulation/policy',\n",
       "  (0.2903772294521332, 'Technical Analysis'),\n",
       "  (0.29612115025520325, 'Economy/Interest Rates'),\n",
       "  (0.29725342988967896, 'Geopolitics'),\n",
       "  (0.29779478907585144, 'Product Development'),\n",
       "  (0.2992168664932251, 'Mergers/Acquisitions'),\n",
       "  (0.30308517813682556, 'Economic Aid'),\n",
       "  (0.3033493757247925, 'Environmental'),\n",
       "  (0.30485522747039795, 'Corruption'),\n",
       "  (0.3065531849861145, 'Safety Concerns'),\n",
       "  (0.3079742193222046, 'Investing'),\n",
       "  (0.30834341049194336, 'Company Performance'),\n",
       "  (0.3132922947406769, 'Economic Outlook'),\n",
       "  (0.3227875828742981, 'Education/Training'),\n",
       "  (0.32524222135543823, 'Economic Data'),\n",
       "  (0.3254878520965576, 'Competition'),\n",
       "  (0.3260963559150696, 'Operations'),\n",
       "  (0.32622283697128296, 'Sales/Revenue'),\n",
       "  (0.32839733362197876, 'Elections'),\n",
       "  (0.3284817934036255, 'Protests/Strikes'),\n",
       "  (0.33366477489471436, 'Labor/Workforce'),\n",
       "  (0.3373696208000183, 'Marketing/Partnerships'),\n",
       "  (0.33804428577423096, 'Utilities'),\n",
       "  (0.3430916666984558, 'Labor Negotiations'),\n",
       "  (0.34377914667129517, 'Airline Industry'),\n",
       "  (0.34802091121673584, 'International Trade'),\n",
       "  (0.35171929001808167, 'European Union'),\n",
       "  (0.3535033166408539, 'Infrastructure/Construction'),\n",
       "  (0.35509681701660156, 'Legal'),\n",
       "  (0.3743923306465149, 'Consumer Goods'),\n",
       "  (0.3754892647266388, 'Politics'),\n",
       "  (0.3816375136375427, 'Taxation'),\n",
       "  (0.3828403651714325, 'Security Measures'),\n",
       "  (0.41155296564102173, 'Financial/Economic'),\n",
       "  (0.42010584473609924, 'Insurance'),\n",
       "  (0.4368091821670532, 'Government/Public Sector'),\n",
       "  (0.45738935470581055, 'Privatization'),\n",
       "  (0.5002402067184448, 'Budget/Fiscal Policy'),\n",
       "  (0.6432530283927917, 'Legislation'),\n",
       "  (0.666974663734436, 'Trade Regulations'),\n",
       "  (0.8851657509803772, 'Regulation')]]"
      ]
     },
     "execution_count": 325,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim_list = []\n",
    "for i, sim_score in enumerate(sim_matrix):\n",
    "    rank_list = np.argsort(sim_score)[-40:]\n",
    "    cur_list = [true_label[i]]\n",
    "    for index in rank_list:\n",
    "        cur_list.append((sim_score[index], predict_label_space[index]))\n",
    "    if cur_list[-1][0] >= 0.6:\n",
    "        sim_list.append(cur_list)\n",
    "sim_list[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 326,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[('sports', 'Sports'), 1.0000001192092896],\n",
       " [('mergers/acquisitions', 'Mergers/Acquisitions'), 1.0],\n",
       " [('international relations', 'International Relations'), 1.0],\n",
       " [('fashion', 'Fashion'), 1.0],\n",
       " [('health', 'Health'), 1.0]]"
      ]
     },
     "execution_count": 326,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pair_list = []\n",
    "for line in sim_list:\n",
    "    true_label = line[0]\n",
    "    for row in line[1:]:\n",
    "        if row[0] >= 0.6:\n",
    "            pair_list.append([(true_label, row[1]), row[0]])\n",
    "sorted_pair_list = sorted(pair_list, key=lambda x: x[1], reverse = True)\n",
    "sorted_pair_list[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 327,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[('sports', 'Sports'), 1.0000001192092896],\n",
       " [('mergers/acquisitions', 'Mergers/Acquisitions'), 1.0],\n",
       " [('international relations', 'International Relations'), 1.0],\n",
       " [('fashion', 'Fashion'), 1.0],\n",
       " [('health', 'Health'), 1.0]]"
      ]
     },
     "execution_count": 327,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "refine_pair_list = []\n",
    "while len(sorted_pair_list) > 0:\n",
    "    refine_pair_list.append([sorted_pair_list[0][0], sorted_pair_list[0][1]])\n",
    "    ground_truth = sorted_pair_list[0][0][0]\n",
    "    prediction = sorted_pair_list[0][0][1]\n",
    "    sorted_pair_list.pop(0)\n",
    "    delete_list = []\n",
    "    for row in sorted_pair_list:\n",
    "        if row[0][0] == ground_truth or row[0][1] == prediction:\n",
    "            delete_list.append(row)\n",
    "    for item in delete_list:\n",
    "        sorted_pair_list.remove(item)\n",
    "refine_pair_list[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 328,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "44"
      ]
     },
     "execution_count": 328,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(refine_pair_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.42718446601941745"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "44/103"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Multi-Label",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
